import argparse
import os
import platform
import shutil
import time
from pathlib import Path
import cv2
import numpy as np

import torch
import torch.backends.cudnn as cudnn

from ultralytics.utils.downloads import attempt_download_asset
from ultralytics.utils.checks import check_imgsz, check_imshow
from ultralytics.utils.torch_utils import select_device, time_sync
from ultralytics.data.augment import LetterBox
from ultralytics.utils.ops import non_max_suppression
from ultralytics.nn.tasks import attempt_load_weights
from deep_sort_pytorch.utils.parser import get_config
from deep_sort_pytorch.deep_sort import DeepSort
from utils import scale_coords

palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1)
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'


def xyxy_to_xywh(*xyxy):
    """" Calculates the relative bounding box from absolute pixel values. """
    bbox_left = min([xyxy[0].item(), xyxy[2].item()])
    bbox_top = min([xyxy[1].item(), xyxy[3].item()])
    bbox_w = abs(xyxy[0].item() - xyxy[2].item())
    bbox_h = abs(xyxy[1].item() - xyxy[3].item())
    x_c = (bbox_left + bbox_w / 2)
    y_c = (bbox_top + bbox_h / 2)
    w = bbox_w
    h = bbox_h
    return x_c, y_c, w, h


def xyxy_to_tlwh(bbox_xyxy):
    tlwh_bboxs = []
    for i, box in enumerate(bbox_xyxy):
        x1, y1, x2, y2 = [int(i) for i in box]
        top = x1
        left = y1
        w = int(x2 - x1)
        h = int(y2 - y1)
        tlwh_obj = [top, left, w, h]
        tlwh_bboxs.append(tlwh_obj)
    return tlwh_bboxs


def compute_color_for_labels(label):
    """
    Simple function that adds fixed color depending on the class
    """
    color = [int((p * (label ** 2 - label + 1)) % 255) for p in palette]
    return tuple(color)


def draw_boxes(img, bbox, identities=None, offset=(0, 0)):
    for i, box in enumerate(bbox):
        x1, y1, x2, y2 = [int(i) for i in box]
        x1 += offset[0]
        x2 += offset[0]
        y1 += offset[1]
        y2 += offset[1]
        # box text and bar
        id = int(identities[i]) if identities is not None else 0
        color = compute_color_for_labels(id)
        label = '{}{:d}'.format("", id)
        t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 2, 2)[0]
        cv2.rectangle(img, (x1, y1), (x2, y2), color, 3)
        cv2.rectangle(
            img, (x1, y1), (x1 + t_size[0] + 3, y1 + t_size[1] + 4), color, -1)
        cv2.putText(img, label, (x1, y1 +
                                 t_size[1] + 4), cv2.FONT_HERSHEY_PLAIN, 2, [255, 255, 255], 2)
    return img


def detect(opt):
    out, source, yolo_weights, deep_sort_weights, show_vid, save_vid, save_txt, imgsz, evaluate = \
        opt.output, opt.source, opt.yolo_weights, opt.deep_sort_weights, opt.show_vid, opt.save_vid, \
            opt.save_txt, opt.img_size, opt.evaluate
    
    # 检查是否为摄像头输入
    is_camera = source.isdigit() or source == '0' or source.startswith('rtsp') or source.startswith('http')
    
    # 初始化DeepSort
    cfg = get_config()
    cfg.merge_from_file(opt.config_deepsort)
    attempt_download_asset(deep_sort_weights, repo='mikel-brostrom/Yolov5_DeepSort_Pytorch')
    deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
                        max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
                        nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP, max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
                        max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET,
                        use_cuda=True)

    # 初始化设备
    device = select_device(opt.device)
    half = device.type != 'cpu'  # 半精度仅在CUDA上支持

    # 创建输出目录
    if not evaluate:
        os.makedirs(out, exist_ok=True)

    # 加载模型
    model = attempt_load_weights(yolo_weights, device=device)  # 加载FP32模型
    stride = int(model.stride.max())  # 模型步长
    imgsz = check_imgsz(imgsz, stride=stride)  # 检查图像尺寸
    if half:
        model.half()  # 转换为FP16

    # 初始化视频捕获
    cap = None
    if is_camera:
        # 处理摄像头输入
        source = int(source) if source.isdigit() else 0
        cap = cv2.VideoCapture(source)
        if not cap.isOpened():
            raise IOError(f"Cannot open camera {source}")
        fps = cap.get(cv2.CAP_PROP_FPS)
        w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    else:
        # 处理视频文件输入
        cap = cv2.VideoCapture(source)
        if not cap.isOpened():
            raise IOError(f"Cannot open video file {source}")
        fps = cap.get(cv2.CAP_PROP_FPS)
        w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # 初始化视频写入器
    vid_writer = None
    if save_vid:
        save_path = os.path.join(out, "camera_output.mp4" if is_camera else os.path.basename(source))
        vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))

    # 初始化文本输出
    txt_path = os.path.join(out, "camera_output.txt" if is_camera else os.path.splitext(os.path.basename(source))[0] + '.txt')
    if save_txt:
        open(txt_path, 'w').close()  # 清空现有文件

    # 获取类别名称
    names = model.module.names if hasattr(model, 'module') else model.names

    # 预热模型
    if device.type != 'cpu':
        model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))

    # 初始化LetterBox
    letterbox = LetterBox(imgsz, auto=True, stride=stride)

    # 主循环
    frame_idx = 0
    t0 = time.time()
    
    while cap.isOpened():
        ret, im0 = cap.read()
        if not ret:
            break
            
        # 预处理图像
        img = letterbox(image=im0)
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR转RGB
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(device)
        img = img.half() if half else img.float()  # uint8转fp16/32
        img /= 255.0  # 0-255转0.0-1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        # 推理
        t1 = time_sync()
        pred = model(img, augment=opt.augment)[0]
        
        # 应用NMS
        pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
        t2 = time_sync()

        # 处理检测结果
        det = pred[0]  # 获取第一个(也是唯一一个)批次的检测结果
        s = f'Frame {frame_idx}: '

        if det is not None and len(det):
            # 将边界框从图像尺寸缩放到原始图像尺寸
            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

            # 打印结果
            for c in det[:, -1].unique():
                n = (det[:, -1] == c).sum()  # 每个类别的检测数量
                s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # 添加到字符串

            xywh_bboxs = []
            confs = []

            # 将检测结果适配到DeepSort输入格式
            for *xyxy, conf, cls in det:
                # 转换为DeepSort格式
                x_c, y_c, bbox_w, bbox_h = xyxy_to_xywh(*xyxy)
                xywh_obj = [x_c, y_c, bbox_w, bbox_h]
                xywh_bboxs.append(xywh_obj)
                confs.append([conf.item()])

            xywhs = torch.Tensor(xywh_bboxs)
            confss = torch.Tensor(confs)

            # 将检测结果传递给DeepSort
            outputs = deepsort.update(xywhs, confss, im0)

            # 绘制边界框用于可视化
            if len(outputs) > 0:
                bbox_xyxy = outputs[:, :4]
                identities = outputs[:, -1]
                draw_boxes(im0, bbox_xyxy, identities)
                # 转换为MOT格式
                tlwh_bboxs = xyxy_to_tlwh(bbox_xyxy)

                # 将MOT兼容结果写入文件
                if save_txt:
                    for j, (tlwh_bbox, output) in enumerate(zip(tlwh_bboxs, outputs)):
                        bbox_top = tlwh_bbox[0]
                        bbox_left = tlwh_bbox[1]
                        bbox_w = tlwh_bbox[2]
                        bbox_h = tlwh_bbox[3]
                        identity = output[-1]
                        with open(txt_path, 'a') as f:
                            f.write(('%g ' * 10 + '\n') % (frame_idx, identity, bbox_top,
                                                           bbox_left, bbox_w, bbox_h, -1, -1, -1,
                                                           -1))  # 标签格式
        else:
            deepsort.increment_ages()

        # 打印时间(推理 + NMS)
        print(f'{s}Done. ({(1E3 * (t2 - t1)):.1f}ms)')

        # 显示结果
        # if show_vid:
        cv2.imshow('Camera Tracking', im0)
        if cv2.waitKey(1) == ord('q'):  # 按q退出
            break

        # 保存带检测结果的视频
        if save_vid and vid_writer is not None:
            vid_writer.write(im0)

        frame_idx += 1

    # 释放资源
    cap.release()
    if vid_writer is not None:
        vid_writer.release()
    cv2.destroyAllWindows()

    print(f'Done. ({time.time() - t0:.3f}s total)')
    if save_txt or save_vid:
        print(f'Results saved to {out}')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--yolo_weights', type=str, default=r'C:\codes\python\deepsort\runs\detect\train\weights\best.pt', help='model.pt path')
    parser.add_argument('--deep_sort_weights', type=str, default=r'D:\GitHub\deepsort_YOLO11_pytorch\deep_sort_pytorch\deep_sort\deep\checkpoint\books\ckpt.t7', help='ckpt.t7 path')
    # 文件/摄像头，0表示摄像头
    parser.add_argument('--source', type=str, default='0', help='输入源（摄像头ID或视频路径）')
    parser.add_argument('--output', type=str, default='inference/output', help='输出文件夹')
    parser.add_argument('--img-size', type=int, default=640, help='推理尺寸（像素）')
    parser.add_argument('--conf-thres', type=float, default=0.4, help='目标置信度阈值')
    parser.add_argument('--iou-thres', type=float, default=0.5, help='NMS的IOU阈值')
    parser.add_argument('--fourcc', type=str, default='mp4v', help='输出视频编解码器')
    parser.add_argument('--device', default='', help='CUDA设备，例如0或0,1,2,3或cpu')
    parser.add_argument('--show-vid', action='store_true', help='显示跟踪视频结果')
    parser.add_argument('--save-vid', default=True, action='store_true', help='保存视频跟踪结果')
    parser.add_argument('--save-txt', action='store_true', help='保存MOT兼容结果到*.txt')
    # 类别0是人，1是自行车，2是汽车...79是烤箱
    parser.add_argument('--classes', nargs='+', type=int, help='按类别过滤：--class 0 或 --class 16 17')
    parser.add_argument('--agnostic-nms', action='store_true', help='类别不可知的NMS')
    parser.add_argument('--augment', action='store_true', help='增强推理')
    parser.add_argument('--evaluate', action='store_true', help='评估模式')
    parser.add_argument("--config_deepsort", type=str, default="deep_sort_pytorch/configs/deep_sort.yaml")
    args = parser.parse_args()
    args.img_size = check_imgsz(args.img_size)

    with torch.no_grad():
        detect(args)